#include <opencv2/opencv.hpp>
#include <iostream>
#include <vector>
#include <string>
#include <filesystem>
namespace fs = std::filesystem;


// 计算亮度系数 k 和偏移值 da
// 计算亮度系数 k 和偏移值 da
std::pair<double, double> calculateBrightness(const cv::Mat& img) {
    // 转换为灰度图
    cv::Mat grayImg;
    cv::cvtColor(img, grayImg, cv::COLOR_BGR2GRAY);
    
    int height = grayImg.rows;
    int width = grayImg.cols;
    int size = height * width;

    // 计算直方图
    int histSize = 256;
    float range[] = {0, 256};
    const float* histRange = {range};
    cv::Mat hist;
    cv::calcHist(&grayImg, 1, 0, cv::Mat(), hist, 1, &histSize, &histRange);

    // 计算偏离均值(128)的程度
    cv::Mat reduceMatrix = cv::Mat::ones(height, width, CV_8U) * 128;
    cv::Mat shiftValue;
    cv::subtract(grayImg, reduceMatrix, shiftValue, cv::noArray(), CV_32S);
    double shiftSum = cv::sum(shiftValue)[0];
    double da = shiftSum / size;

    // 计算平均偏差
    double ma = 0.0;
    for (int i = 0; i < histSize; ++i) {
        ma += std::abs(i - 128 - da) * hist.at<float>(i);
    }
    double m = std::abs(ma / size);

    // 亮度系数
    double k = std::abs(da) / m;
    
    std::cout << "k: " << k << std::endl;
    std::cout << "da: " << da << std::endl;
    
    return std::make_pair(k, da);
}

// 处理单张图片
void processImage(const std::string& imagePath, const std::string& outputDir) {
    cv::Mat img = cv::imread(imagePath);
    if (img.empty()) {
        std::cerr << "无法读取图片: " << imagePath << std::endl;
        return;
    }

    auto result = calculateBrightness(img);
    double k = result.first;
    double da = result.second;

    std::string status;
    cv::Scalar color;

    if (k > 1) {
        if (da > 0) {
            status = "Overexposed";
            color = cv::Scalar(0, 0, 255); // 红色
        } else {
            if (da >= -80) {
                status = "Slightly Dark";
                color = cv::Scalar(0, 165, 255); // 橙色
            } else if (da >= -110) {
                status = "Moderately Dark";
                color = cv::Scalar(0, 0, 255); // 红色
            } else {
                status = "Very Dark";
                color = cv::Scalar(255, 0, 0); // 蓝色
            }
        }
    } else {
        status = "Normal";
        color = cv::Scalar(0, 255, 0); // 绿色
    }

    // 渲染文本
    std::string text1 = "Brightness: " + status;
    std::string text2 = "k = " + std::to_string(k).substr(0, 6) +
                        ", da = " + std::to_string(da).substr(0, 6);

    cv::putText(img, text1, cv::Point(20, 40), cv::FONT_HERSHEY_SIMPLEX, 1, color, 2);
    cv::putText(img, text2, cv::Point(20, 80), cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(255, 255, 255), 2);

    // 保存处理后的图片
    std::string filename = fs::path(imagePath).filename().string();
    std::string outputPath = outputDir + "/" + filename;
    cv::imwrite(outputPath, img);
    std::cout << "已处理: " << filename << " -> " << status << std::endl;
}

// 处理文件夹中的所有图片
void processFolder(const std::string& inputDir, const std::string& outputDir) {
    if (!fs::exists(outputDir)) {
        fs::create_directories(outputDir);
    }

    for (const auto& entry : fs::directory_iterator(inputDir)) {
        std::string ext = entry.path().extension().string();
        if (ext == ".png" || ext == ".jpg" || ext == ".jpeg" || ext == ".bmp") {
            processImage(entry.path().string(), outputDir);
        }
    }
}

int main() {
    std::string input_folder = "/ai_data1/wireope_seg/testData/DA5804532"; // 输入路径
    std::string output_folder = "output_images2"; // 输出路径

    processFolder(input_folder, output_folder);

    std::cout << "所有图片处理完成！" << std::endl;
    return 0;
}
